#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
相似度方法测试

测试SentenceTransformerRetriever的多种相似度计算方法
"""

import sys
import os

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 在导入任何库之前强制设置环境变量
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
os.environ["HUGGINGFACE_CO_URL_HOME"] = "https://hf-mirror.com"
os.environ["HUGGINGFACE_HUB_CACHE"] = os.path.expanduser("~/.cache/huggingface")
os.environ["SENTENCE_TRANSFORMERS_HOME"] = os.path.expanduser("~/.cache/sentence_transformers")

# 强制导入配置，确保镜像环境变量已设置
import questionretrieval.config
from loguru import logger

logger.info("🌐 使用 HuggingFace 国内镜像配置")
logger.info(f"📁 缓存目录: {questionretrieval.config.HF_HUB_CACHE}")
logger.info(f"🔗 镜像地址: {questionretrieval.config.HF_MIRROR_ENDPOINT}")

import pytest
import numpy as np
from questionretrieval import load_questions

# 测试数据
TEST_QUESTIONS = [
    "大伙房水库的主管部门是哪个单位？",
    "大伙房水库总库容/防洪库容/兴利库容/死库容？",
    "大伙房水库调度目标和任务是什么？",
    "大伙房水库防洪调度的核心任务是什么？",
    "大伙房水库供水调度的核心任务是什么？"
]

TEST_QUERY = "大伙房水库的管理和调度"

def get_test_questions():
    """
    获取测试问题
    """
    try:
        questions = load_questions()
        if len(questions) > 0:
            return questions
    except:
        pass
    return TEST_QUESTIONS

def test_sentence_transformer_cosine_similarity():
    """
    测试 SentenceTransformer 余弦相似度方法
    """
    logger.info("\n=== 测试 SentenceTransformer 余弦相似度 ===")
    
    try:
        from questionretrieval import SentenceTransformerRetriever
        retriever = SentenceTransformerRetriever(similarity_method='cosine')
        questions = get_test_questions()
        
        retriever.load_candidates(questions)
        result = retriever.recommend(TEST_QUERY, top_k=3)
        
        if result:
            logger.info(f"✓ 余弦相似度推荐成功，返回 {len(result)} 个结果")
            for i, (question, score, similarity) in enumerate(result[:3], 1):
                logger.info(f"  {i}. {question} (综合分数: {score:.3f}, 相似度: {similarity:.3f})")
            return True
        else:
            logger.warning("✗ 余弦相似度推荐失败，无结果")
            return False
    
    except ImportError:
        logger.warning("✗ SentenceTransformerRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ 余弦相似度测试失败: {e}")
        return False

def test_sentence_transformer_pearson_similarity():
    """
    测试 SentenceTransformer 皮尔逊相关系数相似度方法
    """
    logger.info("\n=== 测试 SentenceTransformer 皮尔逊相关系数相似度 ===")
    
    try:
        from questionretrieval import SentenceTransformerRetriever
        
        retriever = SentenceTransformerRetriever(similarity_method='pearson')
        questions = get_test_questions()
        retriever.load_candidates(questions)
        
        query = "如何学习编程？"
        result = retriever.recommend_questions(query, top_k=3)
        
        if result:
            logger.info(f"✓ 皮尔逊相关系数相似度推荐成功，返回 {len(result)} 个结果")
            for i, (question, score, similarity) in enumerate(result[:3], 1):
                logger.info(f"  {i}. {question} (综合分数: {score:.3f}, 相似度: {similarity:.3f})")
            return True
        else:
            logger.warning("✗ 皮尔逊相关系数相似度推荐失败，无结果")
            return False
    
    except ImportError:
        logger.warning("✗ SentenceTransformerRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ 皮尔逊相关系数相似度测试失败: {e}")
        return False

def test_sentence_transformer_euclidean_similarity():
    """
    测试 SentenceTransformer 欧几里得距离相似度方法
    """
    logger.info("\n=== 测试 SentenceTransformer 欧几里得距离相似度 ===")
    
    try:
        from questionretrieval import SentenceTransformerRetriever
        
        retriever = SentenceTransformerRetriever(similarity_method='euclidean')
        questions = get_test_questions()
        retriever.load_candidates(questions)
        
        query = "如何学习编程？"
        result = retriever.recommend_questions(query, top_k=3)
        
        if result:
            logger.info(f"✓ 欧几里得距离相似度推荐成功，返回 {len(result)} 个结果")
            for i, (question, score, similarity) in enumerate(result[:3], 1):
                logger.info(f"  {i}. {question} (综合分数: {score:.3f}, 相似度: {similarity:.3f})")
            return True
        else:
            logger.warning("✗ 欧几里得距离相似度推荐失败，无结果")
            return False
    
    except ImportError:
        logger.warning("✗ SentenceTransformerRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ 欧几里得距离相似度测试失败: {e}")
        return False

def test_sentence_transformer_manhattan_similarity():
    """
    测试 SentenceTransformer 曼哈顿距离相似度方法
    """
    logger.info("\n=== 测试 SentenceTransformer 曼哈顿距离相似度 ===")
    
    try:
        from questionretrieval import SentenceTransformerRetriever
        
        retriever = SentenceTransformerRetriever(similarity_method='manhattan')
        questions = get_test_questions()
        retriever.load_candidates(questions)
        
        query = "如何学习编程？"
        result = retriever.recommend_questions(query, top_k=3)
        
        if result:
            logger.info(f"✓ 曼哈顿距离相似度推荐成功，返回 {len(result)} 个结果")
            for i, (question, score, similarity) in enumerate(result[:3], 1):
                logger.info(f"  {i}. {question} (综合分数: {score:.3f}, 相似度: {similarity:.3f})")
            return True
        else:
            logger.warning("✗ 曼哈顿距离相似度推荐失败，无结果")
            return False
    
    except ImportError:
        logger.warning("✗ SentenceTransformerRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ 曼哈顿距离相似度测试失败: {e}")
        return False

def test_sentence_transformer_linear_similarity():
    """
    测试 SentenceTransformer 线性核相似度方法
    """
    logger.info("\n=== 测试 SentenceTransformer 线性核相似度 ===")
    
    try:
        from questionretrieval import SentenceTransformerRetriever
        
        retriever = SentenceTransformerRetriever(similarity_method='linear_kernel')
        questions = get_test_questions()
        retriever.load_candidates(questions)
        
        query = "如何学习编程？"
        result = retriever.recommend_questions(query, top_k=3)
        
        if result:
            logger.info(f"✓ 线性核相似度推荐成功，返回 {len(result)} 个结果")
            for i, (question, score, similarity) in enumerate(result[:3], 1):
                logger.info(f"  {i}. {question} (综合分数: {score:.3f}, 相似度: {similarity:.3f})")
            return True
        else:
            logger.warning("✗ 线性核相似度推荐失败，无结果")
            return False
    
    except ImportError:
        logger.warning("✗ SentenceTransformerRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ 线性核相似度测试失败: {e}")
        return False

def test_similarity_method_switching():
    """
    测试相似度方法动态切换功能
    """
    logger.info("\n=== 测试相似度方法动态切换 ===")
    
    try:
        from questionretrieval import SentenceTransformerRetriever
        retriever = SentenceTransformerRetriever()
        questions = get_test_questions()
        retriever.load_candidates(questions)
        
        methods = ['cosine', 'pearson', 'euclidean', 'manhattan', 'linear']
        results = {}
        
        for method in methods:
            try:
                retriever.set_similarity_method(method)
                result = retriever.recommend(TEST_QUERY, top_k=1)
                if result:
                    results[method] = result[0][2]  # 获取相似度分数
                    logger.info(f"  {method}: 相似度 = {result[0][2]:.3f}")
                else:
                    results[method] = 0
                    logger.info(f"  {method}: 无结果")
            except Exception as e:
                logger.error(f"  {method}: 失败 - {e}")
                results[method] = None
        
        success_count = sum(1 for v in results.values() if v is not None)
        logger.info(f"✓ 相似度方法切换测试完成，{success_count}/{len(methods)} 个方法可用")
        return success_count > 0
        
    except ImportError:
        logger.warning("✗ SentenceTransformerRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ 相似度方法切换测试失败: {e}")
        return False

def test_similarity_matrix_calculation():
    """
    测试相似度矩阵计算功能
    """
    logger.info("\n=== 测试相似度矩阵计算 ===")
    
    try:
        from questionretrieval import SentenceTransformerRetriever
        retriever = SentenceTransformerRetriever()
        questions = get_test_questions()[:3]  # 使用较少问题以减少计算时间
        
        retriever.load_candidates(questions)
        similarity_matrix = retriever.get_similarity_matrix()
        
        if similarity_matrix is not None:
            logger.info(f"✓ 相似度矩阵计算成功，形状: {similarity_matrix.shape}")
            logger.info(f"  矩阵对角线元素 (自相似度): {np.diag(similarity_matrix)}")
            logger.info(f"  矩阵最大值: {np.max(similarity_matrix):.3f}")
            logger.info(f"  矩阵最小值: {np.min(similarity_matrix):.3f}")
            return True
        else:
            logger.warning("✗ 相似度矩阵计算失败")
            return False
            
    except ImportError:
        logger.warning("✗ SentenceTransformerRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ 相似度矩阵计算测试失败: {e}")
        return False

def test_all_similarity_methods():
    """
    测试所有相似度方法
    """
    logger.info("\n" + "="*60)
    logger.info("相似度方法测试")
    logger.info("="*60)
    
    results = {
        'cosine': test_sentence_transformer_cosine_similarity(),
        'pearson': test_sentence_transformer_pearson_similarity(),
        'euclidean': test_sentence_transformer_euclidean_similarity(),
        'manhattan': test_sentence_transformer_manhattan_similarity(),
        'linear': test_sentence_transformer_linear_similarity(),
        'method_switching': test_similarity_method_switching(),
        'similarity_matrix': test_similarity_matrix_calculation()
    }
    
    logger.info("\n" + "-"*50)
    logger.info("相似度方法测试结果")
    logger.info("-"*50)
    
    success_count = 0
    for name, success in results.items():
        status = "✓ 成功" if success else "✗ 失败"
        logger.info(f"{name:20} : {status}")
        if success:
            success_count += 1
    
    logger.info(f"\n总计: {success_count}/{len(results)} 个测试通过")
    
    if success_count > 0:
        logger.info("\n💡 提示: SentenceTransformerRetriever 支持多种相似度计算方法")
        logger.info("   可以根据具体应用场景选择最适合的相似度方法")
    else:
        logger.warning("\n⚠️  所有相似度方法测试失败，请检查依赖安装")
        logger.info("解决方案: pip install sentence-transformers scikit-learn")
    
    return results

def test_bert_similarity_methods():
    """
    测试 BERT 相似度方法
    """
    logger.info("\n=== 测试 BERT 相似度方法 ===")
    
    try:
        from questionretrieval import BertSimilarityRetriever
        
        retriever = BertSimilarityRetriever()
        questions = get_test_questions()
        retriever.load_candidates(questions)
        
        query = "如何学习编程？"
        result = retriever.recommend_questions(query, top_k=3)
        
        if result:
            logger.info(f"✓ BERT 相似度推荐成功，返回 {len(result)} 个结果")
            for i, (question, score, similarity) in enumerate(result[:3], 1):
                logger.info(f"  {i}. {question} (综合分数: {score:.3f}, 相似度: {similarity:.3f})")
            return True
        else:
            logger.warning("✗ BERT 相似度推荐失败，无结果")
            return False
    
    except ImportError:
        logger.warning("✗ BertSimilarityRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ BERT 相似度测试失败: {e}")
        return False

if __name__ == "__main__":
    logger.info("\n" + "="*50)
    logger.info("开始测试相似度方法")
    logger.info("="*50)
    
    test_results = []
    
    # 测试 SentenceTransformer 余弦相似度
    test_results.append(test_sentence_transformer_cosine_similarity())
    
    # 测试 SentenceTransformer 皮尔逊相关系数相似度
    test_results.append(test_sentence_transformer_pearson_similarity())
    
    # 测试 SentenceTransformer 欧几里得距离相似度
    test_results.append(test_sentence_transformer_euclidean_similarity())
    
    # 测试 SentenceTransformer 曼哈顿距离相似度
    test_results.append(test_sentence_transformer_manhattan_similarity())
    
    # 测试 SentenceTransformer 线性核相似度
    test_results.append(test_sentence_transformer_linear_kernel_similarity())
    
    # 测试 BERT 相似度方法
    test_results.append(test_bert_similarity_methods())
    
    # 统计结果
    passed = sum(test_results)
    total = len(test_results)
    
    logger.info(f"\n\n=== 相似度方法测试结果 ===")
    logger.info(f"总计: {total} 个测试")
    logger.info(f"通过: {passed} 个")
    logger.info(f"失败: {total - passed} 个")
    
    if passed == total:
        logger.info("\n🎉 所有相似度方法测试通过！")
    else:
        logger.warning("\n⚠️  部分相似度方法测试失败，请检查依赖和配置")
        logger.info("\n💡 提示:")
        logger.info("   - 确保已安装 sentence-transformers")
        logger.info("   - 确保已安装 transformers")
        logger.info("   - 检查网络连接以下载模型")
        logger.info("   - 如果是首次运行，模型下载可能需要时间")